package com.softfly.flink;

import org.apache.flink.api.common.io.DefaultInputSplitAssigner;
import org.apache.flink.api.common.io.RichInputFormat;
import org.apache.flink.api.common.io.statistics.BaseStatistics;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.core.io.InputSplit;
import org.apache.flink.core.io.InputSplitAssigner;
import org.apache.flink.types.Row;
import java.io.IOException;
import java.util.Random;


public class RandomInputFormat extends RichInputFormat<Row, InputSplit> {

    private int current = 0;

    private Random random = new Random();

    @Override
    public void configure(Configuration configuration) {

    }

    @Override
    public BaseStatistics getStatistics(BaseStatistics baseStatistics) throws IOException {
        return null;
    }

    @Override
    public InputSplit[] createInputSplits(int num) throws IOException {
        InputSplit[] inputSplits = new InputSplit[num];
        for (int i = 0; i < inputSplits.length; ++i) {
            inputSplits[i] = new RandomInputSplit(i, num);
        }
        return inputSplits;
    }

    @Override
    public InputSplitAssigner getInputSplitAssigner(InputSplit[] inputSplits) {
        return new DefaultInputSplitAssigner(inputSplits);
    }

    @Override
    public void open(InputSplit inputSplit) throws IOException {
        System.out.println("inputformat open");
    }

    @Override
    public boolean reachedEnd() throws IOException {
//        if (current == 10) {
//            return true;
//        }
//        return false;
        return false;
    }

    @Override
    public Row nextRecord(Row row) throws IOException {
        row = new Row(2);
        row.setField(0, random.nextInt(10));
        row.setField(1, random.nextInt(20));
        current++;
        return row;
    }

    @Override
    public void close() throws IOException {
        System.out.println("RandomInputFormat close()");
    }

}
